# ----------------------------------------------------------------------
# Numenta Platform for Intelligent Computing (NuPIC)
# Copyright (C) 2020, Numenta, Inc.  Unless you have an agreement
# with Numenta, Inc., for a separate license for this software code, the
# following terms and conditions apply:
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero Public License version 3 as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
# See the GNU Affero Public License for more details.
#
# You should have received a copy of the GNU Affero Public License
# along with this program.  If not, see http://www.gnu.org/licenses.
#
# http://numenta.org/licenses/
# ----------------------------------------------------------------------

import json
import logging
from pathlib import Path

import ax
import numpy as np
from ax.modelbridge.generation_strategy import (
    GenerationStep,
    GenerationStrategy,
    Models,
)
from ax.service.ax_client import AxClient
from ray import tune
from ray.tune.suggest.suggestion import SuggestionAlgorithm


def ax_client_with_explicit_strategy(num_random, num_computed):
    steps = []
    if num_random > 0:
        steps.append(GenerationStep(model=Models.SOBOL,
                                    num_arms=num_random))
    if num_computed > 0:
        steps.append(GenerationStep(model=Models.GPEI, num_arms=-1))

    return AxClient(enforce_sequential_optimization=False,
                    generation_strategy=GenerationStrategy(steps))


def get_ray_trials(experiment_dir, parameters):
    trials = []

    experiment_dir = Path(experiment_dir)
    if experiment_dir.exists():
        analysis = tune.Analysis(experiment_dir)
        for _, result in analysis.dataframe().iterrows():
            with open(Path(result["logdir"]) / "params.json", "r") as fin:
                full_config = json.load(fin)

            config = tuple(sorted(
                [(parameter["name"], full_config[parameter["name"]])
                 for parameter in parameters]))

            trials.append((config, result))
    else:
        print(f"Experiment directory {experiment_dir} does not yet exist.")

    return trials


def filter_to_pareto(trials, loss_metrics):

    scores = np.empty((len(trials), len(loss_metrics)),
                      dtype=np.float)

    for i, (config, result) in enumerate(trials):
        for j, metric in enumerate(loss_metrics):
            scores[i, j] = metric(config, result)

    filtered = []

    for i in range(len(trials)):
        s = scores[i]
        on_frontier = True
        for i2 in range(len(trials)):
            if i2 != i:
                s2 = scores[i2]
                if (s2 < s).all():
                    on_frontier = False

        if on_frontier:
            filtered.append(trials[i])

    return filtered


RAY_LOGGER = logging.getLogger(__name__)


class AxSearch(SuggestionAlgorithm):
    """
    Adapted from the Ray Tune source. Rather then returning SEM 0 (as the Ray
    source does), it returns None, allowing Ax to estimate the SEM.

    A wrapper around Ax to provide trial suggestions.
    Requires Ax to be installed. Ax is an open source tool from
    Facebook for configuring and optimizing experiments. More information
    can be found in https://ax.dev/.
    Parameters:
        parameters (list[dict]): Parameters in the experiment search space.
            Required elements in the dictionaries are: "name" (name of
            this parameter, string), "type" (type of the parameter: "range",
            "fixed", or "choice", string), "bounds" for range parameters
            (list of two values, lower bound first), "values" for choice
            parameters (list of values), and "value" for fixed parameters
            (single value).
        objective_name (str): Name of the metric used as objective in this
            experiment. This metric must be present in `raw_data` argument
            to `log_data`. This metric must also be present in the dict
            reported/returned by the Trainable.
        max_concurrent (int): Number of maximum concurrent trials. Defaults
            to 10.
        minimize (bool): Whether this experiment represents a minimization
            problem. Defaults to False.
        parameter_constraints (list[str]): Parameter constraints, such as
            "x3 >= x4" or "x3 + x4 >= 2".
        outcome_constraints (list[str]): Outcome constraints of form
            "metric_name >= bound", like "m1 <= 3."
        use_early_stopped_trials (bool): Whether to use early terminated
            trial results in the optimization process.
    Example:
        >>> parameters = [
        >>>     {"name": "x1", "type": "range", "bounds": [0.0, 1.0]},
        >>>     {"name": "x2", "type": "range", "bounds": [0.0, 1.0]},
        >>> ]
        >>> algo = AxSearch(parameters=parameters,
        >>>     objective_name="hartmann6", max_concurrent=4)
    """

    def __init__(self, ax_client, max_concurrent=10, **kwargs):
        assert ax is not None, "Ax must be installed!"
        assert type(max_concurrent) is int and max_concurrent > 0
        self._ax = ax_client
        exp = self._ax.experiment
        self._objective_name = exp.optimization_config.objective.metric.name
        if self._ax._enforce_sequential_optimization:
            RAY_LOGGER.warning("Detected sequential enforcement. Setting max "
                               "concurrency to 1.")
            max_concurrent = 1
        self._max_concurrent = max_concurrent
        self._parameters = list(exp.parameters)
        self._live_index_mapping = {}
        super(AxSearch, self).__init__(**kwargs)

    def _suggest(self, trial_id):
        if self._num_live_trials() >= self._max_concurrent:
            return None
        parameters, trial_index = self._ax.get_next_trial()
        self._live_index_mapping[trial_id] = trial_index
        return parameters

    def on_trial_result(self, trial_id, result):
        pass

    def on_trial_complete(self,
                          trial_id,
                          result=None,
                          error=False,
                          early_terminated=False):
        """Notification for the completion of trial.
        Data of form key value dictionary of metric names and values.
        """
        if result:
            self._process_result(trial_id, result, early_terminated)
        self._live_index_mapping.pop(trial_id)

    def _process_result(self, trial_id, result, early_terminated=False):
        if early_terminated and self._use_early_stopped is False:
            return
        ax_trial_index = self._live_index_mapping[trial_id]
        metric_dict = {
            self._objective_name: (result[self._objective_name], None)
        }
        outcome_names = [
            oc.metric.name for oc in
            self._ax.experiment.optimization_config.outcome_constraints
        ]
        metric_dict.update({on: (result[on], None) for on in outcome_names})
        self._ax.complete_trial(
            trial_index=ax_trial_index, raw_data=metric_dict)

    def _num_live_trials(self):
        return len(self._live_index_mapping)
